import akshare as ak
import os
import pandas as pd
from datetime import datetime, timedelta
import numpy as np


def get_file_path(name):
    """
    :param name: 文件名
    :type name: str
    :return: 文件地址
    :type: str
    """

    # 当前文件夹
    current_dir = os.path.dirname(__file__)

    # 获取股票数据文件夹
    stock_data_path = os.path.join(os.path.dirname(current_dir), 'stock_data')

    # 数据地址
    file_path = os.path.join(stock_data_path, name)

    return file_path


# 获取交易日历
def get_trading_days(start_date, end_date):
    """获取指定日期范围内的交易日列表"""
    trade_cal = ak.stock_zh_a_daily(symbol="sh000001", start_date=start_date, end_date=end_date)
    trading_days = [d.strftime('%Y%m%d') for d in trade_cal['date'].tolist()]
    return trading_days


# 获取前5天的首次涨停数据
def get_past_5_days_first_zt_data():
    """
    :return: 前五个交易日内涨停一次的股票数据（包含涨停日期）
    :type: pandas.DataFrame
    """
    # 获取当前时间
    today = datetime.now()
    today_str = today.strftime('%Y%m%d')
    print(today)

    # 计算足够大的日期范围以确保包含5个交易日
    start_date = (today - timedelta(days=15)).strftime('%Y%m%d')
    trading_days = get_trading_days(start_date, today_str)  # 交易日列表（按时间升序）

    # 如果今天不是交易日，取最近的前一个交易日
    if today_str in trading_days:
        current_trading_day = today_str
    else:
        current_trading_day = trading_days[-1]
    # 获取当前交易日在交易日历中的索引
    current_trading_idx = trading_days.index(current_trading_day)

    # 取最近的5个交易日
    past_5_trading_days = trading_days[-5:]
    print(past_5_trading_days)
    df = pd.DataFrame()

    # 遍历前5个交易日
    for date in past_5_trading_days:
        date_obj = pd.to_datetime(date)
        date_str = date_obj.strftime('%Y%m%d')

        # 找到涨停日期在交易日历中的索引
        zt_trading_idx = trading_days.index(date_str)
        # 索引差 = 当前交易日索引 - 涨停日期索引
        trading_days_since = current_trading_idx - zt_trading_idx

        # 获取指定日期的涨停数据
        stock_zt_data = ak.stock_zt_pool_em(date=date_str)

        # 添加“距离今天的交易日天数”
        stock_zt_data['距今涨停交易日'] = trading_days_since
        # 添加涨停日期
        stock_zt_data['涨停日期'] = date_str

        # 追加到结果
        df = pd.concat([df, stock_zt_data], ignore_index=True)

    temp_df = df.copy()
    # 删除连板数据
    df_unique = temp_df.drop_duplicates(subset='名称', keep=False)

    # 删除连板数不是1，涨停统计不是1/1
    result_df = df_unique[
        (df_unique['连板数'] == 1) &
        (df_unique['涨停统计'] == '1/1') &
        (df_unique['距今涨停交易日'] != 0)
    ]
    return result_df, temp_df


# 删除st类型股票
def del_st(df):
    """
    :param df: 股票数据
    :type df: pandas.DataFrame
    :return: 删除st类型的股票数据
    :type: pandas.DataFrame
    """
    df = df[~df['名称'].str.contains("ST")]
    return df


# 删除科创板
def del_kc(df):
    """
    :param df: 股票数据
    :type df: pandas.DataFrame
    :return: 删除科创板股票数据
    :type: pandas.DataFrame
    """
    df = df[~df['代码'].str.startswith('688')]
    return df


# 删除创业板
def del_cy(df):
    """
    :param df: 股票数据
    :type df: pandas.DataFrame
    :return: 删除创业板股票数据
    :type: pandas.DataFrame
    """
    df = df[~df['代码'].str.startswith('30')]
    return df


# 删除北证所
def del_bse(df):
    return df[~df['代码'].str.startswith('8')]


# 筛选数据
def select_data(df):
    df = del_st(df)
    df = del_kc(df)
    df = del_cy(df)
    df = del_bse(df)
    return df


# 优化数据
def optimize_data(df):
    """
    :param df: 股票数据
    :type df: pandas.DataFrame
    :return: 优化后的股票数据
    :type: pandas.DataFrame
    """
    # 涨跌幅保留2位小数
    df['涨跌幅'] = df['涨跌幅'].round(2)

    # 成交额单位转为亿元，并且保留2位小数
    df['成交额'] = df['成交额'] / 100000000
    df['成交额（亿元）'] = df['成交额'].round(2)
    del df['成交额']

    # 流通市值单位转为亿元，并且保留2位小数
    df['流通市值'] = df['流通市值'] / 100000000
    df['流通市值（亿元）'] = df['流通市值'].round(2)
    del df['流通市值']

    # 总值单位转为亿元，并且保留2位小数
    df['总市值'] = df['总市值'] / 100000000
    df['总市值（亿元）'] = df['总市值'].round(2)
    del df['总市值']

    # 换手率单位转为 %，并且保留2位小数
    df['换手率'] = df['换手率'] * 100
    df['换手率（%）'] = df['换手率'].round(2)
    del df['换手率']

    # 封板资金单位转为亿 ，并且保留2位小数
    df['封板资金'] = df['封板资金'] / 100000000
    df['封板资金（亿元）'] = df['封板资金'].round(2)
    del df['封板资金']

    # 首次封板时间由时分秒转为时:分:秒
    df['首次封板时间'] = df['首次封板时间'].str[:2]+':'+df['首次封板时间'].str[2:4]+':'+df['首次封板时间'].str[4:]

    # 最后封板时间由时分秒转为时:分:秒
    df['最后封板时间'] = df['最后封板时间'].str[:2]+':'+df['最后封板时间'].str[2:4]+':'+df['最后封板时间'].str[4:]

    # 筛选150亿流通市值以下的股票
    df = df[df['流通市值（亿元）'] < 150]

    return df


# 获取数据
data_past_5_days_first_zt, data_past_5_days_first_zt_temp = get_past_5_days_first_zt_data()
# 筛选数据
data_past_5_days_first_zt = select_data(data_past_5_days_first_zt)
# 优化数据
data_past_5_days_first_zt = optimize_data(data_past_5_days_first_zt)
# 保存数据
data_past_5_days_first_zt.to_csv(get_file_path('首板股票数据.csv'), encoding='utf-8')
# data_past_5_days_first_zt_temp.to_csv(get_file_path('stock_zt_data_temp.csv'), encoding='utf-8')
